
from NodeTree import *

class RedBlackTree:
    def __init__ (self):
        self.emptyNode = NodeTree(None) #Atributo raiz da classe arvore vermelha preta.
        self.emptyNode.setColor('PRETO')
        self.emptyNode.setLeft(self.emptyNode)
        self.emptyNode.setRight(self.emptyNode)
        self.emptyNode.setFather(self.emptyNode)
        self.root = self.emptyNode

    def getRoot(self):
        return  self.root;

    
    #Funcao de inserir.
    def insert(self, dado):
        y = self.emptyNode #Ponteiro
        x = self.root #No temporario
        z = NodeTree(dado) #Novo no
        z.setLeft(self.emptyNode)
        z.setRight(self.emptyNode)
        z.setFather(self.emptyNode)
        
        while (x.getData() != None):
            y = x
            if z.getData() < x.getData():
                x = x.getLeft()#self.root.setLeft(self.emptyNode)
            else:
                x = x.getRight()#self.root.setRight(self.emptyNode)

        z.setFather(y)
        if y.getData() == None:
            self.root = z
        elif z.getData() < y.getData():
            y.setLeft(z)
        else:
            y.setRight(z)
        z.setColor ('VERMELHO')
        
        self.FixUpInsert(z)

    #Rotacionar a esquera.
    def LeftRotate(self, x):
        y = x.getRight() #Define y.
        x.setRight(y.getLeft()) #Faz da subarvore esquerda de y a subarvore direita de x.
        y.getLeft().setFather(x)
        y.setFather(x.getFather()) #Liga o pai de x a y.
        if x.getFather().getData() == None:
            self.root = y
        elif x == x.getFather().getLeft():
            x.getFather().setLeft(y)
        else:
            x.getFather().setRight(y)
        y.setLeft(x) #Coloca x a esquerda de y.
        x.setFather(y)

    #Rotacionar a direita.
    def RightRotate(self, x):
        y = x.getLeft() #Define y.
        x.setLeft(y.getRight()) #Faz da subarvore direita de y a subarvore esquerda de x.
        y.getRight().sefFather(x)
        y.setFather(x.getFather()) #Liga o pai de x a y.
        if x.getFather() == None:
            self.root = y
        elif x == x.getFather().getRight():
            x.getFather().setRightt(y)
        else:
            x.getFather().setLeft(y)
        y.setRight(x) #Coloca x a direita de y.
        x.setFather(y)

    #Funcao de consertar ao inserir.    
    def FixUpInsert (self, z):
        while z.getFather().getColor() == "VERMELHO":
            if z.getFather() == z.getFather().getFather().getLeft():
                y = z.getFather().getFather.getRight()
                if y.getColor() == "VERMELHO":
                    z.getFather().setColor('PRETO')
                    y.setColor('PRETO')
                    z.getFather().getFather().setColor('VERMELHO')
                    z = z.getFather().getFather()
                else: 
                    if z == z.getFather().getRight():
                        z = z.getFather()
                        self.LeftRotate(z)
                    z.getFather().setColor('PRETO')
                    z.getFather().getFather().setColor('VERMELHO')
                    self.RightRotate(z.getFather().getFather())
            else:
                y = z.getFather().getFather().getLeft()
                if y.getColor() == "VERMELHO":
                    z.getFather().setColor('PRETO')
                    y.setColor('PRETO')
                    z.getFather().getFather().setColor('VERMELHO')
                    z = z.getFather().getFather()
                else: 
                    if z == z.getFather().getLeft():
                        z = z.getFather()
                        self.RightRotate(z)
                    z.getFather().setColor('PRETO')
                    z.getFather().getFather().setColor('VERMELHO')
                    self.LeftRotate(z.getFather().getFather())
        self.root.setColor('PRETO')

    #Funcao de deletar.
    def delete (self, dado):
        y = self.emptyNode #Ponteiro
        x = self.root #No temporario
        z = NodeTree(dado) #No a ser removido
        z.setLeft(self.emptyNode)
        z.setRight(self.emptyNode)
        z.setFather(self.emptyNode)
        
        if z.getLeft().getData() == None or z.getRight().getData() == None:
            y = z
        else:
            y = TeeSuccessor(z)
        if y.getLeft().getData() != None:
            x = y.getLeft()
        else:
            x = y.getRight()
        x.setFather(y.getFather())
        if y.getFather() == None:
            self.root = x
        elif y == y.getFather().getLeft():
            y.getLeft().setFather(x)
        else:
            y.getRight().setFather(x)
        if y != z:
            z.setData(y.getData())
        if y.getColor() == 'PRETO':
            self.DeleteFixUp(x)
        return y

    #Funcao de consertar ao deletar.
    def DeleteFixUP (self, x):
        while x != self.root and x.getColor() == 'PRETO':
            if x == x.getFather().getLeft():
                w = x.getFather().getRight()
                if w.getColor() == 'VERMELHO':
                    w.setColor('PRETO')
                    w.getFather().setColor('VERMELHO')
                    self.LeftRotate(x.getFather())
                    w = x.getFather().getRight()
                if w.getLeft().getColor() == 'PRETO' and w.getRight().getColor() == 'PRETO':
                    w.setColor('VERMELHO')
                    x = x.getFather()
                elif w.getRight().getColor() == 'PRETO':
                    w.getLeft().setColor('PRETO')
                    w.setColor('VERMELHO')
                    self.RightRotate(w)
                    w = x.getFather().getRight()
                    w.setColor(x.getFather().getColor())
                    x.getFather().setColor('PRETO')
                    w.getRight().setColor('PRETO')
                    self.LeftRotate(x.getFather())
                    x = self.root
            else:
                w = x.getFather().getLeft()
                if w.getColor() == 'VERMELHO':
                    w.setColor('PRETO')
                    w.getFather().setColor('VERMELHO')
                    self.RightRotate(x.getFather())
                    w = x.getFather().getLeft()
                if w.getRight().getColor() == 'PRETO' and w.getLeft().getColor() == 'PRETO':
                    w.setColor('VERMELHO')
                    x = x.getFather()
                elif w.getLeft().getColor() == 'PRETO':
                    w.getRight().setColor('PRETO')
                    w.setColor('VERMELHO')
                    self.LeftRotate(w)
                    w = x.getFather().getLeft()
                    w.setColor(x.getFather().getColor())
                    x.getFather().setColor('PRETO')
                    w.getLeft().setColor('PRETO')
                    self.RightRotate(x.getFather())
                    x = self.root
        x.setColor('PRETO')

    #Para um nodo x, essa funcao encontra seu sucessor.
    def Successor(self, x):

        if x.getRight().getData() == None:
            return self.Minimun(x.getRight())
        y = x.getFather()
        while y.getData() != None and x == y.getRight():
            x = y
            y = y.getFather()
        return y

    #Chave maxima da arvore.
    def Minimun (self, x):
        while x.getLeft().getData() != None:
            x = x.getLeft()
        return x

    #Pesquisa de elemento sobre uma arvore binaria.
    def InorderTreeWalk(self, x):
        if x.getData() != None:
            self.InorderTreeWalk(x.getLeft())
            print x.getData()
            print x.getColor()
            self.InorderTreeWalk(x.getRight())
